package com.example.backend.aspect;

import com.example.backend.annotation.RateLimit;
import com.example.backend.service.TokenBucketService;
import com.example.backend.util.Result;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import jakarta.servlet.http.HttpServletRequest;
import java.util.Objects;

@Aspect
@Component
public class RateLimitAspect {

    private static final Logger logger = LoggerFactory.getLogger(RateLimitAspect.class);
    private final TokenBucketService tokenBucketService;

    public RateLimitAspect(TokenBucketService tokenBucketService) {
        this.tokenBucketService = tokenBucketService;
        logger.info("RateLimitAspect initialized with TokenBucketService");
    }

    @Around("@annotation(rateLimit)")
    public Object checkRateLimit(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws Throwable {
        // 1. 尝试获取当前请求
        HttpServletRequest request;
        try {
            request = ((ServletRequestAttributes) Objects.requireNonNull(
                    RequestContextHolder.currentRequestAttributes())).getRequest();
        } catch (IllegalStateException e) {
            // 非Web环境直接放行
            logger.warn("RateLimitAspect executed in non-web context, bypassing rate limit");
            return joinPoint.proceed();
        }

        if (request == null) {
            logger.warn("Failed to get HttpServletRequest, bypassing rate limit");
            return joinPoint.proceed();
        }

        // 2. 获取用户标识
        String userKey = request.getHeader("X-User-Header");
        if (userKey == null || userKey.isBlank()) {
            userKey = request.getRemoteAddr(); // 无Header时使用IP
            logger.debug("Using client IP as user key: {}", userKey);
        } else {
            logger.debug("Using X-User-Header as user key: {}", userKey);
        }

        // 3. 检查令牌桶
        if (!tokenBucketService.allowRequest(userKey)) {
            logger.warn("Rate limit exceeded for user: {}", userKey);
            return Result.error("429", "请求过于频繁，请稍后再试");
        }

        logger.debug("Request allowed for user: {}", userKey);

        // 4. 放行请求
        return joinPoint.proceed();
    }
}